package com.lucifer.hawkeye.ai.service.impl;

import com.lucifer.hawkeye.ai.advisor.HawkeyeVectorStoreChatMemoryAdvisor;
import com.lucifer.hawkeye.ai.common.response.RestResult;
import com.lucifer.hawkeye.ai.config.HawkeyeAiConfig;
import com.lucifer.hawkeye.ai.domain.conversation.structured.TopN;
import com.lucifer.hawkeye.ai.domain.conversation.type.ConversationType;
import com.lucifer.hawkeye.ai.model.ChatClientFactory;
import com.lucifer.hawkeye.ai.prompt.HawkeyePrompt;
import com.lucifer.hawkeye.ai.service.ChatApiService;
import com.lucifer.hawkeye.ai.service.Nl2SqlService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;


/**
 * @author lucifer
 * @date 2025/3/29 11:32
 */
@Slf4j
@Service
public class ChatApiServiceImpl implements ChatApiService {

    @Resource
    private HawkeyeAiConfig hawkeyeAiConfig;


    @Resource(name = "docVectorStore")
    private VectorStore docVectorStore;

    @Resource(name = "chatMemoryVectorStore")
    private PgVectorStore chatMemoryVectorStore;


    @Resource
    private Nl2SqlService nl2SqlService;

    @Override
    public RestResult conversation(String conversationId, String chatType, String message) {
        ChatClient chatClient = ChatClientFactory.buildChatClient(hawkeyeAiConfig.chatModel());
        ChatClient.CallResponseSpec call1 = chatClient.prompt(message).call();
        String text = call1.chatResponse().getResult().getOutput().getText();
        return RestResult.buildSuccessResult(text);
    }

    @Override
    public RestResult conversationV1(String conversationId,String conversationType, String chatType, String message) {
        String text = "";
        if(Objects.equals(ConversationType.T2T.getCode(),conversationType)){
            Prompt prompt = prompt(conversationId, message);
            ChatClient.CallResponseSpec call = ChatClient.create(hawkeyeAiConfig.chatModel())
                    .prompt(prompt)
                    .advisors(
                            HawkeyeVectorStoreChatMemoryAdvisor.builder(chatMemoryVectorStore).conversationId(conversationId).chatMemoryRetrieveSize(10).order(0).build()
                    )
                    .call();
             text = call.chatResponse().getResult().getOutput().getText();
        }
        if(Objects.equals(ConversationType.NL2SQL.getCode(),conversationType)){
            Map trainMetadata = nl2SqlService.getTrainMetadata(conversationId);
            String datasource = (String)trainMetadata.get("datasource");
            String tableName = (String)trainMetadata.get("table_name");
            text = nl2SqlService.generateSql(conversationId, message, datasource, tableName);
        }
        return RestResult.buildSuccessResult(text);
    }

    @Override
    public Flux<ChatResponse> conversationV2(String conversationId, String chatType, String message) {
        Prompt prompt = prompt(conversationId, message);
        ChatClient.StreamResponseSpec stream = ChatClient.create(hawkeyeAiConfig.chatModel())
                .prompt(prompt)
                .advisors(
                        HawkeyeVectorStoreChatMemoryAdvisor.builder(chatMemoryVectorStore).conversationId(conversationId).chatMemoryRetrieveSize(10).order(0).build()
                )
                .stream();
        return stream.chatResponse();
    }

    @Override
    public TopN topN(String message, String text) {
        Message systemMessage = HawkeyePrompt.chatHotKeyPromptTemplate().createMessage(Map.of("text", text));
        UserMessage userMessage = new UserMessage(message);
        TopN topN = ChatClient.create(hawkeyeAiConfig.chatModel())
                .prompt(new Prompt(List.of(systemMessage, userMessage)))
                .call()
                .entity(TopN.class);
        return topN;
    }

    private Prompt prompt(String conversationId,String message) {
        Prompt prompt = null;
        if (hawkeyeAiConfig.knowledge) {
            String transformerText = "";
//            if (hawkeyeAiConfig.queryTransformer) {
//                transformerText = queryTransformer(conversationId, message);
//            }
            String knowledgeDocument = knowledgeSearch(message);
            StringBuilder documentsBuilder = new StringBuilder();
            if (!StringUtils.isEmpty(knowledgeDocument)) {
                documentsBuilder.append(knowledgeDocument);
            } else {
                documentsBuilder.append(transformerText);
            }
            // 使用Spring AI 提供的模板方式构建SystemMessage对象
            Message systemMessage = HawkeyePrompt.chatSystemPromptTemplate().createMessage(Map.of("documents", documentsBuilder.toString()));
            UserMessage userMessage = new UserMessage(message);
            prompt = new Prompt(List.of(systemMessage, userMessage));
        }else {
            prompt = new Prompt(message);
        }
        return prompt;
    }


//    private String queryTransformer(String conversationId, String message) {
//        Filter.Expression chatMemoryFilterExpression = new FilterExpressionBuilder().eq("conversationId", conversationId).build();
//        List<Document> listOfMemoryDocuments = chatMemoryVectorStore.similaritySearch(SearchRequest.builder()
//                .query(message)
//                .similarityThresholdAll()
//                .topK(30)
//                .filterExpression(chatMemoryFilterExpression)
//                .build());
//
//        List<Message> history = Lists.newArrayList();
//        Optional.ofNullable(listOfMemoryDocuments).orElse(Lists.newArrayList())
//                .stream()
//                .forEach(document -> {
//                    String text = document.getText();
//                    Map<String, Object> metadata = document.getMetadata();
//                    if (metadata.get("messageType").equals("USER")) {
//                        history.add(new UserMessage(text));
//                    }
//                    if (metadata.get("messageType").equals("ASSISTANT")) {
//                        history.add(new AssistantMessage(text));
//                    }
//                });
//
//        Query query = Query.builder()
//                .text(message)
//                .history(history)
//                .build();
//
//        QueryTransformer queryTransformer = CompressionQueryTransformer.builder()
//                .chatClientBuilder(ChatClient.builder(hawkeyeAiConfig.chatModel()))
//                .promptTemplate(HawkeyePrompt.chatCompressionQueryTransformerPromptTemplate())
//                .build();
//
//
//        Query transformedQuery = queryTransformer.transform(query);
//
//        String text = transformedQuery.text();
//
//        return text;
//    }


    private String knowledgeSearch(String message) {
        // 根据问题文本进行相似性搜索
        FilterExpressionBuilder db = new FilterExpressionBuilder();
        Filter.Expression expression = db.or(db.or(db.eq("name", message), db.eq("keyWordTop1", message)), db.or(db.eq("keyWordTop2", message), db.eq("keyWordTop3", message))).build();
        List<Document> listOfDataDocuments = docVectorStore.similaritySearch(SearchRequest.builder()
                .query(message)
                .filterExpression(expression)
                .similarityThresholdAll()
                .topK(3)
                .build());
        // 将Document列表中每个元素的content内容进行拼接获得documents
        String documents = listOfDataDocuments.stream().map(Document::getText).collect(Collectors.joining());
        return documents;
    }

}
